"""
地形校正TVDI
Tm = To + h * k
k = 0.006℃/m
换算成开尔文：k = 273.156K/m
"""
import os.path

from osgeo import gdal
import numpy as np

# 地表温度校正模型
def get_Tm(T, h):
    # T为校正前温度，h为像元高程值，k为高程对地表温度反演过程的平均影响系数
    k = 0.006
    return T + h * k

def generate_func(dem_path, lst_path, out):
    dem_dataset = gdal.Open(dem_path)
    lst_dataset = gdal.Open(lst_path)
    dem = dem_dataset.ReadAsArray()
    lst = lst_dataset.ReadAsArray()
    dem_geo = dem_dataset.GetGeoTransform()
    lst_geo = lst_dataset.GetGeoTransform()
    dem = dem.astype(np.float32)
    lst = lst.astype(np.float32)
    print("dem_array.shape:", dem.shape)
    print("lst_array.shape:", lst.shape)
    lst_m = np.ones_like(lst) * -9999
    for i in range(lst.shape[0]):
        for j in range(lst.shape[1]):
            if lst[i, j] == -9999:
                continue
            # 根据lst坐标获取经纬度
            lst_x = lst_geo[0] + lst_geo[1] * j + lst_geo[2] * i
            lst_y = lst_geo[3] + lst_geo[4] * j + lst_geo[5] * i
            # 根据经纬度获取dem像素坐标
            dem_x = (lst_x - dem_geo[0]) / dem_geo[1]
            dem_y = (lst_y - dem_geo[3]) / dem_geo[5]
            h = dem[int(dem_y), int(dem_x)]
            lst_m[i, j] = get_Tm(lst[i, j], h)
    lst_m = lst_m.astype(np.float32)
    output_path = out + "/" + lst_path.split("\\")[-1].replace(".tif", "_Tm.tif")
    driver = gdal.GetDriverByName('GTiff')
    out_dataset = driver.Create(output_path, lst.shape[1], lst.shape[0], 1, gdal.GDT_Float32)
    out_dataset.SetGeoTransform(lst_geo)
    out_dataset.SetProjection(lst_dataset.GetProjection())
    out_dataset.GetRasterBand(1).SetNoDataValue(-9999)
    out_dataset.GetRasterBand(1).WriteArray(lst_m)
    out_dataset = None
    print("lst_Tm.tif生成完毕")


if __name__ == '__main__':
    dem = r"G:\test\TSVI\dem_clip.tif"
    lst = r"G:\test\TSVI\LST"
    out = r"G:\test\TSVI\LST_Tm"
    if not os.path.exists(out):
        os.makedirs(out)
    lst_list = [os.path.join(lst, i) for i in os.listdir(lst)]
    for lst in lst_list:
        generate_func(dem, lst, out)